"""
用于处理sdf数据
先做二维fft，然后选取固定频率的波形，
之后再做逆fft
"""
import os
import matplotlib as mpl
import sdf_helper as sh
import matplotlib.pyplot as plt
import numpy as np
import analyse as al
from analyse import sim_parm as smp
#==========自定义函数=============================================
#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
# #=========绘图参数==============================
fig = plt.figure(figsize=(22, 10))
#============开始处理=================================================================
#====读取数据==========================================
print('loading data %s%d' % (smp.sdf_prefix, smp.sdf_scope[0]))
data = sh.getdata('%s%.4d.sdf' % (smp.sdf_prefix, smp.sdf_scope[0]))
var = data.Electric_Field_Ey
#2d
X, Y = np.meshgrid(var.grid_mid.data[0][smp.xslice],
                   var.grid_mid.data[1][smp.yslice])
Z = var.data.T[smp.yslice, smp.xslice]

#====test range end==========================================
[rows, cols] = X.shape
print("[rows, cols] =", rows, cols)
Z = np.zeros([rows, cols])
for i in range(rows):
    for j in range(cols):
        # Z[i, j] = al.mathfunc.spherical_wave(0.4e-12,X[i,j],Y[i,j],0 )
        Z[i, j] = al.mathfunc.spherical_wave(0, X[i, j], Y[i, j], 0)
#=================================================================
print("ploting init wave: wavenumber=", al.lp.wavenumber)
ax = plt.subplot(1, 2, 1)
tempax = al.plot.plot_2d(Z, X, Y, ax, ifsdf=0)
print(Z.shape)
#==================================================================
#========================1d=======================================
# Z = Z[int((al.sp.yy - al.sp.yx) / 2), :]
# var_fft = np.fft.fft(Z)
# print('len of var_fft', len(var_fft))
# #注释np.fft.fftfreq返回值f形式为，注意和实际的k之间相差2pi
# #f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)   if n is even 共n项
# #f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)   if n is odd  共n项
# var_x_fre = np.fft.fftfreq(smp.y - smp.x, X[0][1] - X[0][0])  #1167136
# # var_x_fre=np.zeros(len(var_fft))
# # for i in range(len(var_fft)):
# # 	var_x_fre[i]=i/(len(var_fft)*(X[0][1] - X[0][0]))

# print('max fre: ', var_x_fre[np.argmax(var_fft)])
# # plt.scatter(var_x_fre, var_fft)
# # plt.show()
# # print("test==", var.grid_mid.data[1][1] - var.grid_mid.data[1][0],
# #       Y[1][0] - Y[0][0])
# print("begin filter")
# allow_error = 0.05  #滤波允许偏离中心频率的范围
# for i in range(al.sp.y - al.sp.x):
#     if (np.absolute(abs(var_x_fre[i]) - al.lp.wavenumber))>np.absolute(allow_error * (al.lp.wavenumber)):
#         var_fft[i] = 0
# var_ifft = np.real(np.fft.ifft(var_fft))
# temp = np.zeros([rows, cols])
# for i in range(rows):
#     for j in range(cols):
#         temp[i, j] = var_ifft[j]
# ax = plt.subplot(1, 2, 2)
# al.plot.plot_2d(temp, X, Y, ax, ifsdf=0)

#==================================================================
#========================2d=======================================
var_fft2 = np.fft.fft2(Z)
print('len of var_fft2', len(var_fft2[1]), len(var_fft2[:, 1]))
#注释np.fft.fftfreq返回值f形式为，注意和实际的k之间相差2pi
#f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)   if n is even 共n项
#f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)   if n is odd  共n项
var_x_fre = np.fft.fftfreq(smp.y - smp.x, X[0][1] - X[0][0])
var_y_fre = np.fft.fftfreq(smp.yy - smp.yx, Y[1][0] - Y[0][0])
# print("test==", var.grid_mid.data[1][1] - var.grid_mid.data[1][0],
#       Y[1][0] - Y[0][0])
allow_error = 0.05  #滤波允许偏离中心频率的范围
for i in range(al.sp.y - al.sp.x):
    for j in range(al.sp.yy - al.sp.yx):
        if (np.absolute(
                np.sqrt(var_x_fre[i]**2 + var_y_fre[j]**2) -
                al.lp.wavenumber)) > (allow_error * (al.lp.wavenumber)):
            var_fft2[j, i] = 0
var_ifft2 = np.real(np.fft.ifft2(var_fft2))
ax = plt.subplot(1, 2, 2)
al.plot.plot_2d(var_ifft2, X, Y, ax, ifsdf=0)
#&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&
#test

# 绘制colorbar
#前面三个子图的总宽度 为 全部宽度的 0.9；剩下的0.1用来放置colorbar
fig.subplots_adjust(right=0.9)
#colorbar 左 下 宽 高
l = 0.92
b = 0.12
w = 0.015
h = 1 - 2 * b

#对应 l,b,w,h；设置colorbar位置；
rect = [l, b, w, h]
cbar_ax = fig.add_axes(rect)
cb = plt.colorbar(tempax, cax=cbar_ax)

print('savefig')
plt.savefig('test')
plt.close()